import torch.nn as nn

from cnn_train.CnnStockConfig import CnnStockConfig


class CnnStockModel(nn.Module):

    def __init__(self):
        super(CnnStockModel, self).__init__()
        self.features = nn.Sequential(
            nn.Conv1d(1, 64, kernel_size=3, padding=1),  # 输出: (32, 30)
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.MaxPool1d(2),

            nn.Conv1d(64, 128, kernel_size=3, padding=1),  # 输出: (64, 15)
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.MaxPool1d(3),

            nn.Conv1d(128, 256, kernel_size=3, padding=1),  # 输出: (128, 7)
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.AdaptiveAvgPool1d(3)
        )

        self.classifier = nn.Sequential(
            nn.Linear(256 * 3, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, CnnStockConfig.num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x
